/************************************************************/
/*    Standard Multinomial Logit Model with Sparse Input    */ 
/*              Author:  Wayne Zhang                        */
/*            actuary_zhang@hotmail.com                     */
/************************************************************/

/**
 * @file nrmm.c
 * @brief Standard Multinomial Logit Model
 * @author Wayne Zhang                         
 */

#include "rmm.h"
#include "utilities.h"
#include "nrmm.h"


cholmod_common c;

SEXP init(){
  M_R_cholmod_start(&c);
  c.final_ll = 1;	    /* LL' form of simplicial factorization */
  return R_NilValue;
}
SEXP finish(){
  M_cholmod_finish(&c);
  return R_NilValue;
}


/**
 * Multiplication of a sparse matrix (X) and a real vector (b)
 *
 * @param X input sparse matrix 
 * @param b the real vector 
 * @param trans 0: Xb; 1: X'b
 * @param out  result vector 
 *
 */
void sdmult_mv(SEXP X, double *b, int trans, double *out){
  int *dim = SLOT_INT_NULL(X, "Dim"), 
    *xp = SP_SLOT(X), 
    *xi = SI_SLOT(X);
  int nO = dim[0], nB = dim[1]; 
  double *xx = SX_SLOT(X);

  if (trans){ /* X'b */
    AZERO(out, nB);
    for (int j = 0; j < nB; j++){
      for (int p = xp[j]; p < xp[j + 1]; p++)
	out[j] += xx[p] * b[xi[p]];
    } 
  } else {
    AZERO(out, nO);
    for (int j = 0; j < nB; j++){
      for (int p = xp[j]; p < xp[j + 1]; p++)
	out[xi[p]] += xx[p] * b[j];
    }
  }
}

/**
 * Cross-product of a sparse matrix
 *
 * @param X input sparse matrix 
 * @param trans 0: X'X; 1 XX'
 *
 * @return a CHM_SP pointer to a sparse matrix which
 * needs to be freed subsequently 
 * 
 */

CHM_SP scrossprod(CHM_SP x, int trans){
  CHM_SP xt, out;
  if (!trans) 
    xt = M_cholmod_transpose(x, x->xtype, &c);
  out = M_cholmod_aat((!trans) ? xt : x, (int *) NULL, 0, x->xtype, &c);
  if(!out) 
    error("scrossprod(): error return from cholmod_aat()");
  // M_cholmod_band_inplace(out, 0, out->ncol, out->xtype, &c);
  //out->stype = 1; //fixme?
  if (!trans) 
    M_cholmod_free_sparse(&xt, &c);
  return out;
}


/**
 * Update the vector of linear predictors given a sparse 
 * design matrix and a vector of coefficients
 *
 * @param X input sparse design matrix 
 * @param beta the vector of coefficients
 * @param eta the vector of linear predictors
 *
 */
void update_eta(SEXP X, double *beta, double *eta){
  sdmult_mv(X, beta, 0, eta);
}

/**
 * Update the estimated probabilities for the multinomial model
 * given the linear predictor 
 *
 * @param eta the vector of linear predictors
 * @param prob the vector of probabilities
 * @param cset a binary vector indicating the (fixed) choice set
 * @param nP the number of purchase occasions
 * @param nC the number of alternatives
 *
 */
void update_prob(double *eta, double *prob, int *gid, int *nalt, int nP){
  double meta = 0, seta = 0;
  for (int j = 0; j < nP; j++){
    meta = dmax(eta, nalt[j]);                           // meta:= max_k (eta)
    for (int i = 0; i < nalt[j]; i++)
      prob[i] = exp(eta[i] - meta);                      // prob:= e^(eta - meta) -- avoid overflow
    seta = dcumsum(prob, nalt[j]);                       // seta:= sum_k prob
    for (int i = 0; i < nalt[j]; i++)
      prob[i] /= seta;                                   // prob:= prob/seta
    prob += nalt[j];                                     // move to the next individual
    eta += nalt[j];
  }  
}


/**
 * Update the estimated probabilities for the multinomial model
 *
 * @param ans an object of class "mm"
 * @param beta if not null, use its value to update the probabilities;
 * if null, use the beta slot of ans to update the probabilities. 
 * 
 */
void rmm_update_prob(SEXP ans, double *beta){
  int *dims = DIMS_SLOT(ans), *nalt = NALT_SLOT(ans),
    *gid = GID_SLOT(ans);
  int nB = dims[nB_POS], nP = dims[nP_POS], 
    nC = dims[nC_POS];
  double *eta = ETA_SLOT(ans), *mu = MU_SLOT(ans);
  if (beta == NULL)
    beta = BETA_SLOT(ans);
  update_eta(X_SLOT(ans), beta, eta);     
  update_prob(eta, mu, gid, nalt, nP);
 
}

/**
 * Update the *NEGATIVE* loglikelihood 
 *
 * @param ans an object of class "mm"
 *
 */
void rmm_update_llik(SEXP ans){
  int *dims = DIMS_SLOT(ans),
    *yi = YI_SLOT(ans);  
  double *mu = MU_SLOT(ans), *ctrl = CTRL_SLOT(ans),
    *wts = WTS_SLOT(ans);
  double l = 0;
  for (int i = 0; i < dims[nP_POS]; i++)
    l += (mu[yi[i]] > 0) ? (wts[yi[i]] * log(mu[yi[i]])) : 0;       // consideration set
  ctrl[llik_POS] = -l;
}


/**
 * Compute the (negative) Hessian matrix of an nrmm object. The resultant 
 * Hessian is stored in the hess slot. 
 *
 * @param ans an object of class "nrmm"
 *
 */
void rmm_update_hess(SEXP ans){
  int *dims = DIMS_SLOT(ans), *gid = GID_SLOT(ans),
    *nalt = NALT_SLOT(ans);
  int nO = dims[nO_POS], nB = dims[nB_POS],
    nP = dims[nP_POS], nC = dims[nC_POS];
  double *mu = MU_SLOT(ans), *wts = WTS_SLOT(ans),
    *hess = HESS_SLOT(ans);
  CHM_SP sX = AS_CHM_SP(X_SLOT(ans)), tPX,  g, h;
  double *xx_old = Calloc(sX->nzmax, double);

  /* The hessian is calculated as hess = h - g
     1) h = X' diag(mu * w) X
     2) g = X' W^{1/2} P  P' W^{1/2} X  where P=bdiag(mu) is a block-diagonal matrix (nO * nP)
  */
  AZERO(hess, nB * nB);  

  /* 0) This creates the block-diagonal matrix P' = sP */
  CHM_SP sP = M_cholmod_allocate_sparse(nP, nO, nO, TRUE, TRUE, 0, 
					CHOLMOD_REAL, &c);
  sP->itype = CHOLMOD_INT;     
  sP->dtype = CHOLMOD_DOUBLE; 
  double *sx = sP->x; 
  int *si = sP->i, *sp = sP->p;
  for (int i = 0; i < nP; i++){                    // specify the i slot
    for (int j = 0; j < nalt[i]; j++)
      si[gid[i] + j] = i; 
  }
  for (int i = 0; i < nO; i++)                     // specify the x slot
    sx[i] = mu[i] * sqrt(wts[i]);
  sp[0] = 0;
  for (int j = 0; j < nO; j++)                     // specify the p slot
    sp[j + 1] = sp[j] + 1;
  
  // 1) Update g
  tPX = M_cholmod_ssmult(sP, sX, 0, CHOLMOD_REAL,  // tPX = P'X
		       /*out sorted:*/ 1, &c);
  
  g = scrossprod(tPX, 0);                          // g := X' P P' X
  int *gi = g->i, *gp = g->p;
  double *gx = g->x;
  for (int j = 0; j < nB; j++){
    for (int p = gp[j]; p < gp[j + 1]; p++)
      hess[j * nB + gi[p]] -= gx[p];               // hess = -g
  } 
  
  // 2) Update h
  int *xi = sX->i, *xp = sX->p;
  double *xx = sX->x;
  Memcpy(xx_old, xx, sX->nzmax);
  for (int j = 0; j < nB; j++){
    for (int p = xp[j]; p < xp[j + 1]; p++)
      xx[p] *= sqrt(mu[xi[p]] * wts[xi[p]]) ;      // sX = diag(sqrt(mu * w)) * X
  } 
  
  h = scrossprod(sX, 0);                           // h = X' diag(mu) X
  int *hi = h->i, *hp = h->p;
  double *hx = h->x;
  for (int j = 0; j < nB; j++){
    for (int p = hp[j]; p < hp[j + 1]; p++)
      hess[j * nB + hi[p]] += hx[p] ;              // hess = -g + h
  } 
  Memcpy(xx, xx_old, sX->nzmax);                   // X is changed above!

  /*  
  SEXP out;
  PROTECT(out = allocVector(VECSXP, 5)) ;                          
  SET_VECTOR_ELT(out, 0, M_chm_sparse_to_SEXP(sX, 0 , 0, 0, "", R_NilValue));
  SET_VECTOR_ELT(out, 1, M_chm_sparse_to_SEXP(sP, 1 , 0, 0, "", R_NilValue));
  SET_VECTOR_ELT(out, 2, M_chm_sparse_to_SEXP(tPX, 1 , 0, 0, "", R_NilValue));
  SET_VECTOR_ELT(out, 3, M_chm_sparse_to_SEXP(g, 1 , 0, 0, "", R_NilValue));
  SET_VECTOR_ELT(out, 4, M_chm_sparse_to_SEXP(h, 1 , 0, 0, "", R_NilValue));
  UNPROTECT(1);
  */

  Free(xx_old);
  M_cholmod_free_sparse(&sP, &c);
  M_cholmod_free_sparse(&tPX, &c);
  M_cholmod_free_sparse(&g, &c);
  M_cholmod_free_sparse(&h, &c);
}


/**
 * Estimate the (non-mixture) multinomial logit model 
 * using the Majorize-Minimization method
 *
 * @param ans an object of class "nrmm"
 * @param beta if not null, use its value to update the probabilities;
 * if null, use the beta slot of ans to update the probabilities. 
 *
 */

void do_nrm_mm(SEXP ans, double *beta){
  int *dims = DIMS_SLOT(ans);
  int nO = dims[nO_POS], nB = dims[nB_POS];
  double *y = Y_SLOT(ans), 
    *eta = ETA_SLOT(ans), *mu = MU_SLOT(ans), 
    *r = RESID_SLOT(ans), *ctrl = CTRL_SLOT(ans);
  if (beta == NULL) beta = BETA_SLOT(ans);
  double *beta_old = Alloca(nB, double),
    *theta_old = Alloca(nB, double),
    *theta_new = Alloca(nB, double);
  R_CheckStack();
  double t = 0;
  SEXP X = X_SLOT(ans);  
  CHM_DN sol, 
        ctheta = N_AS_CHM_DN(theta_new, nB, 1);
  CHM_FR L = L_SLOT(ans);
  R_CheckStack();
  ctrl[cvg_POS] = 0;
  double *wts = WTS_SLOT(ans);
  double w = dmax(wts, nO);

  Memcpy(theta_old, beta, nB);
  
  for (int iter = 0; iter < ctrl[MMit_POS]; iter++){
    R_CheckUserInterrupt();
    rmm_update_prob(ans, beta);
    for (int i = 0; i < nO; i++)	
      r[i] = mu[i] * (1 - mu[i]); 
    t = 2 * dmax(r, nO) * w;                                // t := 2 * max(p * (1 - p))

    for (int i = 0; i < nO; i++)
      r[i] = wts[i] * (y[i] - mu[i])/t + eta[i];            // r := (y - mu)/t + eta, working response
    sdmult_mv(X, r, 1, theta_new);                          // theta_new := X'r            
    
    sol = M_cholmod_solve(CHOLMOD_A, L, ctheta, &c);        // solve theta_new := L^{-t} * L^{-1} * X'r  
    Memcpy(theta_new, (double*)(sol->x), nB);
    M_cholmod_free_dense(&sol, &c);

    t = (iter + 1.0)/(iter + 4.0);                      // apply the Nesterov step
    Memcpy(beta, theta_new, nB);                        // beta := theta_new + l/(l + 3) * (theta_new - theta_old)
    for (int j = 0; j < nB; j++)
      beta[j] += t * (theta_new[j] - theta_old[j]);
    if (dist(beta, beta_old, nB, 1) < ctrl[eps_POS]) 
      break; 
    Memcpy(beta_old, beta, nB);
    Memcpy(theta_old, theta_new, nB);    
  }

  if (dist(beta, beta_old, nB, 1) >= ctrl[eps_POS])          // test to see if max step limit is reached
    ctrl[cvg_POS] = 1;
  
  rmm_update_hess(ans);                                      // update the Hessian matrix
  
}


/**
 * Estimate the (non-mixture) multinomial logit model using the Newton-Raphson method
 * If ctrl[gam_POS] > 0, then this estimates a ridge-type multinomial logit
 * model with a universal penlaty specified by ctrl[gam_POS]
 *
 * @param ans an object of class "nrmm"
 * @param beta if not null, use its value to update the probabilities;
 * if null, use the beta slot of ans to update the probabilities. 
 *
 */

void do_nrm_nr(SEXP ans, double *beta){
  int *dims = DIMS_SLOT(ans);
  int nO = dims[nO_POS], nB = dims[nB_POS];
  double *y = Y_SLOT(ans), 
    *eta = ETA_SLOT(ans), *mu = MU_SLOT(ans), 
    *r = RESID_SLOT(ans), *ctrl = CTRL_SLOT(ans),
    *wts = WTS_SLOT(ans), *hess = HESS_SLOT(ans);
  if (beta == NULL) beta = BETA_SLOT(ans); 
  double *beta_old = Alloca(nB, double),
    *dbeta = Alloca(nB, double),
    *grad = Alloca(nB, double),
    *ihess = Calloc(nB * nB, double);
  R_CheckStack();
  SEXP X = X_SLOT(ans);
  double llik_old, llik_new, step;
  double gamma = ctrl[gam_POS];
  ctrl[cvg_POS] = 0;

  Memcpy(beta_old, beta, nB);
  rmm_update_prob(ans, beta);                               
  rmm_update_llik(ans);
  llik_old = ctrl[llik_POS] + 0.5 * gamma * sqr_length(beta_old, nB);

  for (int iter = 0; iter < RMM_NRITER; iter++){
    R_CheckUserInterrupt();
    // 1) update the gradient g := X'W(y - mu)
    for (int i = 0; i < nO; i++)
      r[i] = (y[i] - mu[i]) * wts[i];                 
    sdmult_mv(X, r, 1, grad);                      
    if (gamma > 0){                                      // for ridge penalty
      for (int i = 0; i < nB; i++)
	grad[i] -= gamma * beta_old[i];
    }
    // 2) update the hessian
    rmm_update_hess(ans);
    if (gamma > 0){                                      // for ridge penalty
      for (int i = 0; i < nB; i++)
	hess[i + i * nB] += gamma;
    }
    // 3) Newton's direction: -H^{-1}g
    solve_po(nB, hess, ihess);
    mult_mv("N", nB, nB, ihess, grad, dbeta);
    for (step = 1; step > RMM_SMIN; step /= 2) {         // step halving 
      for (int j = 0; j < nB; j++) 
	beta[j] = beta_old[j] + step * dbeta[j];
      rmm_update_prob(ans, beta);                             
      rmm_update_llik(ans);      
      llik_new = ctrl[llik_POS] + 0.5 * gamma * sqr_length(beta, nB);
      if (llik_new <= llik_old) {
	llik_old = llik_new;
	break;
      }       
    }
    if (dist(beta, beta_old, nB, 1) < RMM_EPS) 
      break; 
    Memcpy(beta_old, beta, nB);
  }

  if (dist(beta, beta_old, nB, 1) >= RMM_EPS)            // test to see if max step limit is reached
    ctrl[cvg_POS] = 1;
  
  rmm_update_hess(ans);                                  // update the Hessian matrix
  Free(ihess);
}


/**
 * Perform the E-step in the EM algorithm. That is, it computes the posterior mean
 * probability of each subject belonging to each latent class. This function is 
 * used in all types of models in the rmm package, as the E-step is the same across
 * different models, while the M-step differs. 
 *
 * @param ans an object of class "mm"
 * @param beta the vector of coefficients to be used in updating the probabilities. 
 * @param pi the vector of class distributions to be used.
 * @param W a nI by nS vector that stores the posterior for each subject
 * 
 *
 */
void rmm_update_estep(SEXP ans, double *beta, double *pi, double *W){
  int *dims = DIMS_SLOT(ans),  *yi = YI_SLOT(ans),
    *sub = SUB_SLOT(ans);  
  int nB = dims[nB_POS], nP = dims[nP_POS],
    nS = dims[nS_POS], nI = dims[nI_POS];
  double *mu = MU_SLOT(ans);
  if (beta == NULL) beta = BETA_SLOT(ans);
  if (pi == NULL) pi = PI_SLOT(ans);
  double *lpi = Alloca(nS, double);
  R_CheckStack();
  double d, ws;
  int pos;

  AZERO(W, nI * nS);  
  for (int k = 0; k < nS; k++)                        // lpi := log(pi)
    lpi[k] = log(pi[k]);

  for (int k = 0; k < nS; k++){
    rmm_update_prob(ans, beta + nB * k);              // update prob using beta from the kth segment
    for (int i = 0; i < nP; i++)                      // llik for the ith individual
      W[sub[yi[i]] + nI * k] += (mu[yi[i]] > 0) ? log(mu[yi[i]]) : 0; 
    for (int i = 0; i < nI; i++)                      // llik + log(pi)
      W[i + nI * k] += lpi[k];
  }
    
  for (int i = 0; i < nI; i++){
    d = W[i];                                         // find the max of each row
    for (int k = 0; k < nS; k++){
      if (W[i + nI * k] > d)
	d = W[i + nI * k];
    }
    ws = 0;
    for (int k = 0; k < nS; k++){                     // W_ij = exp( W_ij - max_j W_ij)
      pos = i + nI * k;
      W[pos] = exp(W[pos] - d);
      ws += W[pos];                                   // ws is the row sum
    }
    for (int k = 0; k < nS; k++){
      W[i + nI * k] /= ws;
    }
  }
}

/**
 * Estimate the non-regularized multinomial mixture model. 
 * This function also handles when the number of latent segment is one. 
 *
 * @param ans an object of class "nrmm"
 *
 */
SEXP R_do_nrmm(SEXP ans){
  int *dims = DIMS_SLOT(ans), *sub = SUB_SLOT(ans);  
  int nO = dims[nO_POS], nB = dims[nB_POS],
    nP = dims[nP_POS], nC = dims[nC_POS],
    nG = dims[nG_POS], nS = dims[nS_POS],
    nI = dims[nI_POS];
  int nBS = nB * nS;  //length of the beta vector
  double *beta = BETA_SLOT(ans), *pi = PI_SLOT(ans),
    *wts = WTS_SLOT(ans), *ctrl = CTRL_SLOT(ans);
  double *W = Calloc(nI * nS, double),                   // nI by nS matrix of posteriors 
    *beta_old = Alloca(nBS, double);
  R_CheckStack();
  int iter;

  if (nS == 1){ /* standard mnl model */
    if (ctrl[new_POS]) do_nrm_nr(ans, RMM_DNULL);
    else do_nrm_mm(ans, RMM_DNULL);    
  } else {      /* mixture model */
    Memcpy(beta_old, beta, nBS);
    for (iter = 0; iter < ctrl[EMit_POS]; iter++){
      // E-step 
      rmm_update_estep(ans, RMM_DNULL, RMM_DNULL, W);
      
      // M-step 
      // 1) update pi                                     // pi = colMeans(W)
      sum_mat(nI, nS, 2, W, pi);
      for (int k = 0; k < nS; k++)
	pi[k] /= nI;

      // 2) update beta
      for (int k = 0; k < nS; k++){
	for (int i = 0; i < nO; i++)
	  wts[i] = W[sub[i] + k * nI];
	if (ctrl[new_POS]) do_nrm_nr(ans, beta + k * nB);
	else do_nrm_mm(ans, beta + k * nB);        
      }
      
      if (dist(beta, beta_old, nBS, 1) < ctrl[eps_POS]) 
	break; 
      Memcpy(beta_old, beta, nBS);
    }
    if (iter == ctrl[EMit_POS]) 
      ctrl[cvg_POS] = 1;
  }

  Free(W);
  return R_NilValue;
}




/*=======================================
 *    R callable utility functions      * 
 =======================================*/
SEXP R_rmm_update_prob(SEXP ans){
  rmm_update_prob(ans, RMM_DNULL);
  return R_NilValue;
}

SEXP R_rmm_update_hess(SEXP ans){
  rmm_update_hess(ans);
  return R_NilValue;
}

SEXP R_rmm_update_llik(SEXP ans){
  rmm_update_llik(ans);
  return R_NilValue;
}

SEXP R_rmm_update_estep(SEXP ans, SEXP wlam){
  int *dims = DIMS_SLOT(ans);  
  int nS = dims[nS_POS], nI = dims[nI_POS], 
    nB = dims[nB_POS];
  int wl = LENGTH(wlam) ? (INTEGER(wlam)[0] - 1) : 0;
  double *beta = BETA_SLOT(ans), *pi = PI_SLOT(ans);
  SEXP W = allocVector(REALSXP, nI * nS);
  PROTECT(W);
  rmm_update_estep(ans, beta + wl * nB * nS, pi + wl * nS, REAL(W));
  UNPROTECT(1);
  return W;
}
